package com.usoft.thrift.client.pool;

import com.usoft.thrift.client.provider.ThriftServerAddressProvider;
import com.usoft.thrift.exception.ThriftException;
import org.apache.commons.pool2.BasePooledObjectFactory;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.impl.DefaultPooledObject;
import org.apache.thrift.TServiceClient;
import org.apache.thrift.TServiceClientFactory;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TFramedTransport;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.InetSocketAddress;

/**
 * BasePooledObjectFactory
 * 连接池,thrift-client for spring
 */
public class ThriftPooledClientFactory extends BasePooledObjectFactory<TServiceClient> {

    private Logger logger = LoggerFactory.getLogger(getClass());

    private final ThriftServerAddressProvider serverAddressProvider;
    private final TServiceClientFactory<TServiceClient> clientFactory;
    private PoolOperationCallBack callback;

    protected ThriftPooledClientFactory(
            ThriftServerAddressProvider addressProvider,
            TServiceClientFactory<TServiceClient> clientFactory) throws Exception {
        this.serverAddressProvider = addressProvider;
        this.clientFactory = clientFactory;
    }

    public ThriftPooledClientFactory(ThriftServerAddressProvider addressProvider,
                                     TServiceClientFactory<TServiceClient> clientFactory,
                                     PoolOperationCallBack callback) throws Exception {
        this.serverAddressProvider = addressProvider;
        this.clientFactory = clientFactory;
        this.callback = callback;
    }

    @Override
    public TServiceClient create() throws Exception {
        InetSocketAddress address = serverAddressProvider.selector();
        if (address == null) {
            throw new ThriftException("No provider available for remote service");
        }
        TSocket tsocket = new TSocket(address.getHostName(), address.getPort());
        TTransport transport = new TFramedTransport(tsocket);
        TProtocol protocol = new TBinaryProtocol(transport);
        TServiceClient client = this.clientFactory.getClient(protocol);
        transport.open();
        if (callback != null) {
            try {
                callback.create(client);
            } catch (Exception e) {
                logger.warn("create object:{}", e);
            }
        }
        return client;
    }

    @Override
    public PooledObject<TServiceClient> wrap(TServiceClient tServiceClient) {
        return new DefaultPooledObject<TServiceClient>(tServiceClient);
    }

    @Override
    public void destroyObject(PooledObject<TServiceClient> p) throws Exception {
        TServiceClient client = p.getObject();
        if (callback != null) {
            try {
                callback.destroy(client);
            } catch (Exception e) {
                logger.warn("destroyObject:{}", e);
            }
        }
        logger.info("destroyObject:{}", client);
        TTransport pin = client.getInputProtocol().getTransport();
        pin.close();
        TTransport pout = client.getOutputProtocol().getTransport();
        pout.close();
    }

    @Override
    public boolean validateObject(PooledObject<TServiceClient> p) {
        TServiceClient client = p.getObject();
        TTransport pin = client.getInputProtocol().getTransport();
        logger.info("validateObject input:{}", pin.isOpen());
        TTransport pout = client.getOutputProtocol().getTransport();
        logger.info("validateObject output:{}", pout.isOpen());
        return pin.isOpen() && pout.isOpen();
    }
}
